% perform_allocation.m
% 核心内存分配函数，执行地址分配和SPILL决策。
% (版本 29: 修正了 adj_list/predecessors_list 未同步扩展的bug)

function [solution, graph] = perform_allocation(initial_schedule, graph, code_map, l1_cap, ub_cap, l0a_cap, l0b_cap, l0c_cap, debug_target_step, debug_window)

    % --- 1. 初始化 ---
    solution = struct('final_schedule', [], 'memory_offsets', containers.Map('KeyType','int32','ValueType','any'), 'spill_map', containers.Map('KeyType','int32','ValueType','any'), 'total_spill_volume', 0);
    
    node_template = struct(...
        'id', [], 'op', '', 'buf_id', [], 'size', [], ...
        'type', '', 'pipe', '', 'cycles', [], 'bufs', [], ...
        'op_code', 0, 'type_code', 0);

    next_node_id = graph.num_nodes + 1;
    
    cache_managers = containers.Map('KeyType', 'int32', 'ValueType', 'any');
    cache_managers(code_map.L1) = cm_create(l1_cap);
    cache_managers(code_map.UB) = cm_create(ub_cap);
    cache_managers(code_map.L0A) = cm_create(l0a_cap);
    cache_managers(code_map.L0B) = cm_create(l0b_cap);
    cache_managers(code_map.L0C) = cm_create(l0c_cap);
    
    active_buffers = containers.Map('KeyType', 'int32', 'ValueType', 'any');
    spilled_buffers = containers.Map('KeyType', 'int32', 'ValueType', 'any');

    copy_in_bufs = [];
    for i = 1:graph.num_nodes
        if isempty(graph.nodes(i).id)
            continue;
        end
        node = graph.nodes(i);
        if strcmp(node.op, 'COPY_IN') && ~isempty(node.bufs)
            copy_in_bufs = [copy_in_bufs, node.bufs];
        end
    end
    copy_in_bufs = unique(copy_in_bufs);

    % --- 2. 主循环 ---
    step = 1;
    schedule = initial_schedule;
    initial_schedule_length = length(schedule);
    while step <= length(schedule)
        % 安全检查：如果调度长度超过初始值的5倍，跳出循环
        if length(schedule) > initial_schedule_length * 5
            fprintf('警告: 调度长度已超过初始值的5倍 (%d -> %d)，停止分配以避免无限循环。\n', initial_schedule_length, length(schedule));
            break;
        end
        
        node_id = schedule(step);
        node = graph.nodes(node_id);

        is_alloc_request = (node.op_code == code_map.ALLOC);
        is_spill_in_request = strcmp(node.op, 'SPILL_IN');

        if is_alloc_request || is_spill_in_request
            buf_id_to_alloc = node.buf_id;
            if isempty(buf_id_to_alloc)
                step = step + 1; continue;
            end
            
            if is_alloc_request
                if isempty(node.size) || node.size <= 0 || node.type_code == 0
                    step = step + 1; continue;
                end
                size_needed = node.size;
                type_code = node.type_code;
                if isKey(active_buffers, buf_id_to_alloc)
                    error('Allocator 逻辑错误: 在步骤 %d, 节点 %d, 尝试重复分配一个仍在使用的 BufID: %d', step, node_id, buf_id_to_alloc);
                end
            else % is_spill_in_request
                if ~isKey(spilled_buffers, buf_id_to_alloc)
                    step = step + 1; continue;
                end
                spill_info = spilled_buffers(buf_id_to_alloc);
                type_code = spill_info(1);
                size_needed = spill_info(2);
            end

            manager = cache_managers(type_code);
            [manager, offset] = cm_allocate(manager, size_needed, buf_id_to_alloc);
            cache_managers(type_code) = manager;

            if offset ~= -1
                if is_alloc_request
                    solution.memory_offsets(buf_id_to_alloc) = offset;
                else % is_spill_in_request
                    solution.spill_map(node_id) = offset;
                    remove(spilled_buffers, buf_id_to_alloc);
                end
                active_buffers(buf_id_to_alloc) = [type_code, offset, size_needed];
                step = step + 1;
                continue;
            
            else % Spill logic
                candidate_info_map = containers.Map('KeyType', 'int32', 'ValueType', 'any');
                active_keys = cell2mat(keys(active_buffers));
                for k_idx = 1:length(active_keys)
                    candidate_buf_id = active_keys(k_idx);
                    buffer_info = active_buffers(candidate_buf_id);
                    candidate_type_code = buffer_info(1);
                    candidate_size = buffer_info(3);
                    if candidate_type_code ~= type_code, continue; end
                    farthest_use = 0;
                    for future_step = (step + 1):length(schedule)
                        future_node = graph.nodes(schedule(future_step));
                        found = false;
                        if ~isempty(future_node.bufs) && ismember(candidate_buf_id, future_node.bufs)
                            found = true;
                        end
                        if ~found && ~isempty(future_node.buf_id) && future_node.buf_id == candidate_buf_id
                            found = true;
                        end
                        if found
                            farthest_use = future_step;
                            break;
                        end
                    end
                    if farthest_use == 0, farthest_use = length(schedule) + 2; end
                    candidate_info_map(candidate_buf_id) = {farthest_use, candidate_size};
                end

                manager = cache_managers(type_code);
                anchor_points = unique([[manager.free_blocks.start], [manager.allocated_blocks.start]]);
                
                all_plans = [];
                for anchor_addr = anchor_points
                    window_start = anchor_addr;
                    window_end = window_start + size_needed;
                    if window_end > manager.capacity, continue; end
                    
                    current_plan = struct('victim_ids', [], 'cost_future_use', inf, 'cost_volume', 0, 'start_address', anchor_addr);
                    for blk_idx = 1:length(manager.allocated_blocks)
                        alloc_block = manager.allocated_blocks(blk_idx);
                        if alloc_block.start < window_end && (alloc_block.start + alloc_block.size) > window_start
                            if isKey(candidate_info_map, alloc_block.buf_id)
                                current_plan.victim_ids = [current_plan.victim_ids, alloc_block.buf_id];
                            end
                        end
                    end
                    
                    if ~isempty(current_plan.victim_ids)
                        min_future_use = inf;
                        total_volume = 0;
                        for v_id = current_plan.victim_ids
                            info = candidate_info_map(v_id);
                            min_future_use = min(min_future_use, info{1});
                            total_volume = total_volume + info{2};
                        end
                        current_plan.cost_future_use = min_future_use;
                        current_plan.cost_volume = total_volume;
                        all_plans = [all_plans; current_plan];
                    end
                end

                if isempty(all_plans)
                    error('内存分配失败: 无法生成任何有效的Spill方案。');
                end

                T = struct2table(all_plans);
                T = sortrows(T, {'cost_future_use', 'cost_volume'}, {'descend', 'ascend'});
                best_plan = table2struct(T(1,:));
                
                insertions = [];
                for victim_buf_id = best_plan.victim_ids
                    victim_info = active_buffers(victim_buf_id);
                    victim_type_code = victim_info(1);
                    victim_size = victim_info(3);
                    
                    if ismember(victim_buf_id, copy_in_bufs)
                        solution.total_spill_volume = solution.total_spill_volume + victim_size;
                    else
                        solution.total_spill_volume = solution.total_spill_volume + victim_size * 2;
                    end
                    
                    spilled_buffers(victim_buf_id) = [victim_type_code, victim_size];
                    remove(active_buffers, victim_buf_id);
                    
                    manager = cache_managers(victim_type_code);
                    manager = cm_free(manager, victim_buf_id);
                    cache_managers(victim_type_code) = manager;

                    spill_out_node = node_template;
                    spill_out_node.id = next_node_id;
                    spill_out_node.op = 'SPILL_OUT';
                    spill_out_node.bufs = victim_buf_id;
                    spill_out_node.buf_id = victim_buf_id;
                    
                    spill_in_node = node_template;
                    spill_in_node.id = next_node_id + 1;
                    spill_in_node.op = 'SPILL_IN';
                    spill_in_node.bufs = victim_buf_id;
                    spill_in_node.buf_id = victim_buf_id;

                    spill_out_node.pipe = 'MTE3';
                    spill_in_node.pipe = 'MTE2';
                    if ismember(victim_buf_id, copy_in_bufs)
                        spill_out_node.cycles = 0;
                        spill_in_node.cycles = victim_size * 2 + 150;
                    else
                        spill_out_node.cycles = victim_size * 2 + 150;
                        spill_in_node.cycles = victim_size * 2 + 150;
                    end
                    
                    next_node_id = next_node_id + 2;

                    graph.nodes(spill_out_node.id) = spill_out_node;
                    graph.nodes(spill_in_node.id) = spill_in_node;
                    graph.num_nodes = next_node_id - 1;

                    % *** 已修正: 同步扩展邻接表和前驱列表的大小 ***
                    new_size = graph.num_nodes;
                    if new_size > length(graph.adj_list)
                        % 将元胞数组扩展到新大小，并用空元胞初始化新位置
                        graph.adj_list{new_size} = []; 
                        graph.predecessors_list{new_size} = [];
                    end
                    
                    info = candidate_info_map(victim_buf_id);
                    farthest_use = info{1};
                    insertion_point = farthest_use;
                    if insertion_point > length(schedule), insertion_point = length(schedule) + 1; end
                    
                    insertions = [insertions; step, spill_out_node.id];
                    insertions = [insertions; insertion_point, spill_in_node.id];
                end
                
                insertions = sortrows(insertions, 1, 'descend');
                for ins_idx = 1:size(insertions, 1)
                    point = insertions(ins_idx, 1);
                    id_to_insert = insertions(ins_idx, 2);
                    schedule = [schedule(1:point-1); id_to_insert; schedule(point:end)];
                end
            end

        elseif node.op_code == code_map.FREE
            if ~isempty(node.buf_id) && isKey(active_buffers, node.buf_id)
                buf_id = node.buf_id;
                buffer_info = active_buffers(buf_id);
                type_code = buffer_info(1);
                
                manager = cache_managers(type_code);
                manager = cm_free(manager, buf_id);
                cache_managers(type_code) = manager;
                
                remove(active_buffers, buf_id);
            end
            step = step + 1;
        else
            step = step + 1;
        end
    end

    solution.final_schedule = schedule;
end